package njust.csc.aecore.analysis

import io.github.oshai.kotlinlogging.KotlinLogging
import njust.csc.aecore.entity.APK
import njust.csc.aecore.utils.ImageUtils
import njust.csc.aecore.utils.PathUtils
import njust.csc.aecore.utils.RegexUtils
import org.apdplat.word.WordSegmenter
import org.apdplat.word.segmentation.Word
import org.w3c.dom.Element
import org.w3c.dom.NodeList
import java.io.FileNotFoundException
import java.nio.file.Path
import javax.xml.parsers.DocumentBuilderFactory
import kotlin.io.path.*

private val logger = KotlinLogging.logger {}

/**
 * @Author ZhouHang
 * @Date 24/6/24 周一
 */
class DataProcess {

    /**
     * 解析 PCAP 文件，主要是 TLS 中的 SNI 信息
     * */
    fun parsePCAP(apk: APK): Set<String> {
        val pcapFile = PathUtils.getPCAPFile(apk.packageName)
        return parsePCAP(pcapFile)

//        val pcap = Pcap.openStream(pcapFile)
//        pcap.loop { packet ->
//            if (packet.hasProtocol(Protocol.TLS)) {
//                val tlsPacket = packet.getPacket(Protocol.TLS)
//                println(tlsPacket.toString())
//            }
//            true
//        }
    }

    /**
     * @param pcapPath pcap 文件路径
     * */
    fun parsePCAP(pcapPath: Path): Set<String> {
        val bytes = pcapPath.readText()
        val urls = RegexUtils.urlMatch.findAll(bytes)
        val res = mutableSetOf<String>()
        urls.forEach { url ->
            res.add(url.value)
        }
        return res
    }

    /**
     * 处理 fastbot 生成的截图，根据 apk 找到对应的 screencap 位置
     * */
    @Deprecated(message = "失效")
    fun parseScreencap(apk: APK) {
//        val screencapDir = PathUtils.getScreencapDir(apk.packageName)
//        val words = mutableListOf<String>()
//        for (png in screencapDir.listDirectoryEntries()) {
//            words.add(parseScreencap(png))
//        }
//        apk.screenContent = words.joinToString(separator = " ")
    }

    /**
     * 处理 fastbot 生成的截图
     * @param screencapPath 截图文件夹路径
     * @param apkPackageName apk 包名
     * */
    fun parseScreencap(screencapPath: Path, apkPackageName: String): String {
        val allTextSet = mutableSetOf<String>()
        val allContentDescSet = mutableSetOf<String>()
        // 删除多余的 *.log 和 crash 文件夹，节省空间
        for (file in screencapPath.listDirectoryEntries()) {
            if (file.name.contains("Crash")) {
                PathUtils.deleteDir(file)
            } else if (file.extension == "log") {
                file.deleteIfExists()
            }
        }

        // 开始过滤截图
        for (xml in PathUtils.xmlsInSequence(screencapPath)) {
            val png = xml.parent.resolve("${xml.nameWithoutExtension}.png")
            val factory = DocumentBuilderFactory.newInstance()
            val builder = factory.newDocumentBuilder()
            val document = builder.parse(xml.toFile())
            val root = document.documentElement
            val nodeList: NodeList = root.getElementsByTagName("node")
            val textSet = mutableSetOf<String>()
            val contentDescSet = mutableSetOf<String>()

            // 若该截图不是当前分析 app 的截图，则删除
            if (root.getAttribute("package") != apkPackageName) {
                logger.debug { "${png.name} 由于不属于当前app被删除" }
                png.deleteIfExists()
                xml.deleteIfExists()
                continue
            }
            for (i in 0 until nodeList.length) {
                val node = nodeList.item(i) as Element
                val text = node.getAttribute("text")
                val contentDesc = node.getAttribute("content-desc")
                val packageName = node.getAttribute("package")
                val bound = node.getAttribute("bounds")
                if (packageName != apkPackageName)  continue

                if (!text.isNullOrBlank()) textSet.add(text)
                if (!contentDesc.isNullOrBlank()) contentDescSet.add(contentDesc)
            }
            allTextSet.addAll(textSet)
            allContentDescSet.addAll(contentDescSet)

            // 若该图片大面积空白则删除
            var blankRate: Double
            try {
                blankRate = ImageUtils.imageFilter(png)
            } catch (e: FileNotFoundException) {
                continue
            }
            if (blankRate > 0.8) {
                logger.debug { "${png.name} 由于空白太多被删除" }
                png.deleteIfExists()
                xml.deleteIfExists()
            }
        }

        // 计算两张图片的相似度，若相似则删除其中一张
        val xmls = PathUtils.xmlsInSequence(screencapPath)
        for (i in xmls.indices) {
            for (j in i + 1 until xmls.size) {
                val currentXml = xmls[i]
                val nextXml = xmls[j]
                val currentPng = currentXml.parent.resolve("${currentXml.nameWithoutExtension}.png")
                val nextPng = nextXml.parent.resolve("${nextXml.nameWithoutExtension}.png")
                var similarity: Double
                try {
                    similarity = ImageUtils.calculateSimilarity(currentPng, nextPng)
                } catch (e: FileNotFoundException) {
                    continue  // 循环中可能会删掉一些图片，导致这里找不到
                }
                if (similarity > 0.65) {
                    logger.debug { "${nextPng.name} 由于存在相似被删除" }
                    nextXml.deleteIfExists()
                    nextPng.deleteIfExists()
                }
            }
        }
        allTextSet.addAll(allContentDescSet)
        return allTextSet.joinToString(separator = " ")
    }

    /**
     * 将给定字符串分词生成字符串列表
     * @param sentence 待分词的句子
     * @return 经过去重的列表
     * */
    fun wordSeg(sentence: String): List<String> {
        val ret = WordSegmenter.seg(sentence)
        return ret.map { it.text }.distinct()
    }
}